package toolkits

import (
	"fmt"
	"io/ioutil"
	"strconv"
	"strings"

	"github.com/toolkits/pkg/file"
)

type ProcStatus struct {
	Name  string
	State string
	Pid   int
	PPid  int
}

func GetChildPids(pid int, pss []ProcStatus) (childPids []int) {
	childPids = make([]int, 0)

	for _, ps := range pss {
		if ps.PPid == pid {
			childPids = append(childPids, ps.Pid)
		}
	}

	childCount := len(childPids)
	for i := 0; i < childCount; i++ {
		childPid := childPids[i]
		grandsonPids := GetChildPids(childPid, pss)
		childPids = append(childPids, grandsonPids...)
	}

	return
}

func GetAllProcStatus() (pss []ProcStatus, err error) {
	pss = make([]ProcStatus, 0)

	dirs, e := file.DirsUnder("/proc")
	if e != nil {
		err = e
		return
	}

	for i := 0; i < len(dirs); i++ {
		pid, e := strconv.Atoi(dirs[i])
		if e != nil {
			continue
		}

		statusFile := fmt.Sprintf("/proc/%d/status", pid)
		if !file.IsExist(statusFile) {
			continue
		}

		ps, e := ParseProcStatus(statusFile)
		if e != nil {
			continue
		}
		if ps.Pid == 0 {
			continue
		}
		if ps.PPid == 0 {
			continue
		}

		pss = append(pss, ps)
	}

	return
}

func ParseProcStatus(path string) (ps ProcStatus, err error) {
	content, e := ioutil.ReadFile(path)
	if e != nil {
		err = e
		return
	}

	procInfo := make(map[string]string)
	lines := strings.Split(string(content), "\n")
	for _, line := range lines {
		fields := strings.SplitN(line, ":", 2)
		if len(fields) < 2 {
			continue
		}
		key := fields[0]
		value := strings.TrimSpace(fields[1])
		procInfo[key] = value
	}

	{
		ps.Name = procInfo["Name"]
		ps.State = procInfo["State"]
	}

	{
		pids := procInfo["Pid"]
		pid, _ := strconv.Atoi(pids)
		ps.Pid = pid
	}

	{
		ppids := procInfo["PPid"]
		ppid, _ := strconv.Atoi(ppids)
		ps.PPid = ppid
	}

	return
}
